{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "02d4d6f7-1b36-40ed-b26e-a035174c693f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from typing import List\n",
    "\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
    "from langchain_community.graphs import Neo4jGraph\n",
    "\n",
    "os.environ[\"GOOGLE_API_KEY\"] = \"\"\n",
    "\n",
    "llm = ChatGoogleGenerativeAI(model=\"gemini-1.5-pro\", timeout=60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8f260509-ec97-46b5-9a0c-2562afdfb416",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_types = {\n",
    "    \"Simple Retrieval Queries\": \"These queries focus on basic data extraction, retrieving nodes or relationships based on straightforward criteria such as labels, properties, or direct relationships. Examples include fetching all nodes labeled as 'Person' or retrieving relationships of a specific type like 'EMPLOYED_BY'. Simple retrieval is essential for initial data inspections and basic reporting tasks. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Complex Retrieval Queries\": \"These advanced queries use the rich pattern-matching capabilities of Cypher to handle multiple node types and relationship patterns. They involve sophisticated filtering conditions and logical operations to extract nuanced insights from interconnected data points. An example could be finding all 'Person' nodes who work in a 'Department' with over 50 employees and have at least one 'REPORTS_TO' relationship. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Simple Aggregation Queries\": \"Simple aggregation involves calculating basic statistical metrics over properties of nodes or relationships, such as counting the number of nodes, averaging property values, or determining maximum and minimum values. These queries summarize data characteristics and support quick analytical conclusions. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Pathfinding Queries\": \"Specialized in exploring connections between nodes, these queries are used to find the shortest path, identify all paths up to a certain length, or explore possible routes within a network. They are essential for applications in network analysis, routing, logistics, and social network exploration. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Complex Aggregation Queries\": \"The most sophisticated category, these queries involve multiple aggregation functions and often group results over complex subgraphs. They calculate metrics like average number of reports per manager or total sales volume through a network, supporting strategic decision making and advanced reporting. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Verbose query\": \"These queries are characterized by their explicit and detailed specifications about the data retrieval process and the exact information needed. They involve elaborate instructions for navigating through complex data structures, specifying precise criteria for inclusion, exclusion, and sorting of data points. Verbose queries typically require the breakdown of each step in the querying process, from the initial identification of relevant data nodes and relationships to the intricate filtering and sorting mechanisms that must be applied. Always limit the number of results if more than one row is expected from the questions by saying 'first 3' or 'top 5' elements\",\n",
    "    \"Evaluation query\": \"This query type focuses on retrieving specific pieces of data from complex databases with precision. Use clear and detailed instructions to extract relevant information, such as movie titles, product names, or employee IDs, depending on the context. Always ask for a single property or item, titled intuitively based on the data retrieved (e.g., Movie Titles Featuring Tom Cruise). Limit the results to a specific number like 'first 3' or 'top 5' to keep the output concise and focused.\",\n",
    "    \"Multi-step Queries\": \"Multistep queries in a graph database involve executing several operations or traversals to derive the answer. These queries typically combine different data elements by following multiple relationships and filtering nodes at various steps to reach a final result. They often require joining data from various parts of the schema, aggregating results, or applying multiple conditions to uncover complex insights that are not immediately apparent from a single node or relationship\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eb6c23d4-f2c6-4b88-90d3-5a4ef27ada33",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"Your task is to generate 100 questions that are directly related to a specific graph schema in Neo4j. Each question should target distinct aspects of the schema, such as relationships between nodes, properties of nodes, or characteristics of node types. Ensure that the questions vary in complexity, covering basic, intermediate, and advanced queries.\n",
    "Imagine you are a user at a company that needs to present all the types of questions that the graph can answer.\n",
    "You have to be very diligent at your job. Make sure you will accomplish a diversity of questions, ranging from various complexities.\n",
    "\n",
    "Avoid ambiguous questions. For clarity, an ambiguous question is one that can be interpreted in multiple ways or does not have a straightforward answer based on the schema. For example, avoid asking, \"What is related to this?\" without specifying the node type or relationship.\n",
    "Please design each question to yield a limited number of results, specifically between 3 to 10 results. This will ensure that the queries are precise and suitable for detailed analysis and training.\n",
    "The goal of these questions is to create a dataset for training AI models to convert natural language queries into Cypher queries effectively.\n",
    "It is vital that the database contains information that can answer the question!\n",
    "Never write any assumptions, just the questions!!!\n",
    "Make sure to generate 100 questions!\n",
    "\n",
    "Make sure to create questions for the following graph schema:{input}\\n \n",
    "Here are some example nodes and relationship values: {values}. \n",
    "Don't use any values that aren't found in the schema or in provided values.\n",
    "{query_type}\n",
    "Also, do not ask questions that there is no way to answer based on the schema or provided example values. \n",
    "Find good questions that will test the capabilities of graph answering.\n",
    "The output of the should be 1 question per row. Example output format:\n",
    "What movies did Tom Cruise acted in?\n",
    "Which product made the most revenue?\n",
    "Who is the manager of the team that completed the most projects last year?\n",
    "Generated questions:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "935bbf60-d1aa-4738-9db9-9dea319597da",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.prompt import PromptTemplate\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"values\", \"query_type\"], template=prompt_template\n",
    ")\n",
    "\n",
    "chain = prompt | llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "84b30844-54b2-4503-af4f-aed36f8dba2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEMO_URL = \"neo4j+s://demo.neo4jlabs.com\"\n",
    "DEMO_DATABASES = [\n",
    "    \"recommendations\",\n",
    "    \"buzzoverflow\",\n",
    "    \"bluesky\",\n",
    "    \"companies\",\n",
    "    \"fincen\",\n",
    "    \"gameofthrones\",\n",
    "    \"grandstack\",\n",
    "    \"movies\",\n",
    "    \"neoflix\",\n",
    "    \"network\",\n",
    "    \"northwind\",\n",
    "    \"offshoreleaks\",\n",
    "    \"stackoverflow2\",\n",
    "    \"twitch\",\n",
    "    \"twitter\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "131c364d-3983-4db2-84a2-507b3d6b1cf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def remove_enumeration(text):\n",
    "    # This regular expression matches numbers followed by a dot and an optional space at the start of a string\n",
    "    return re.sub(r'^\\d+\\.\\s?', '', text).strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ca56095-ed11-4eb5-b0fb-ffd2cf816877",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_questions = []\n",
    "for database in DEMO_DATABASES:\n",
    "    print(database)\n",
    "    graph = Neo4jGraph(\n",
    "        url=DEMO_URL,\n",
    "        database=database,\n",
    "        username=database,\n",
    "        password=database,\n",
    "        enhanced_schema=True,\n",
    "        sanitize=True,\n",
    "        timeout=30,\n",
    "    )\n",
    "    schema = graph.schema\n",
    "    for type in query_types:\n",
    "        print(type)\n",
    "        instructions = f\"{type}: {query_types[type]}\"\n",
    "        # Sample values\n",
    "        values = graph.query(\n",
    "                \"\"\"MATCH (n)\n",
    "    WHERE rand() > 0.6\n",
    "    WITH n LIMIT 2\n",
    "    CALL {\n",
    "        WITH n\n",
    "        MATCH p=(n)-[*3..3]-()\n",
    "        RETURN p LIMIT 1\n",
    "    }\n",
    "    RETURN p\"\"\"\n",
    "            )\n",
    "\n",
    "        try: # sometimes it timeouts\n",
    "            questions = chain.invoke(\n",
    "                {\"input\": schema, \"query_type\": instructions, \"values\": values}\n",
    "            )\n",
    "            all_questions.extend(\n",
    "            [\n",
    "                {\"question\": remove_enumeration(el), \"type\": type, \"database\": database}\n",
    "                for el in questions.content.split(\"\\n\") if not \"## 100\" in el and el\n",
    "            ]\n",
    "            )\n",
    "        except:\n",
    "            continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3e7a4b77-5669-4dd0-b963-4e4af5d4b660",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>type</th>\n",
       "      <th>database</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Which movies from France have an IMDb rating h...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What are the titles of the top 5 highest-gross...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>List the first 3 movies directed by Harold Lloyd.</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What genres are associated with the movie \"Toy...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What is the average IMDb rating of movies rele...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9962</th>\n",
       "      <td>Find the users who have retweeted tweets from ...</td>\n",
       "      <td>Multi-step Queries</td>\n",
       "      <td>twitter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9963</th>\n",
       "      <td>List the tweets posted by \"Neo4j\" that have be...</td>\n",
       "      <td>Multi-step Queries</td>\n",
       "      <td>twitter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9964</th>\n",
       "      <td>What is the average number of followers for us...</td>\n",
       "      <td>Multi-step Queries</td>\n",
       "      <td>twitter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9965</th>\n",
       "      <td>Which users have a similarity score greater th...</td>\n",
       "      <td>Multi-step Queries</td>\n",
       "      <td>twitter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9966</th>\n",
       "      <td>Find the users who follow \"Neo4j\" and have men...</td>\n",
       "      <td>Multi-step Queries</td>\n",
       "      <td>twitter</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>9967 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               question  \\\n",
       "0     Which movies from France have an IMDb rating h...   \n",
       "1     What are the titles of the top 5 highest-gross...   \n",
       "2     List the first 3 movies directed by Harold Lloyd.   \n",
       "3     What genres are associated with the movie \"Toy...   \n",
       "4     What is the average IMDb rating of movies rele...   \n",
       "...                                                 ...   \n",
       "9962  Find the users who have retweeted tweets from ...   \n",
       "9963  List the tweets posted by \"Neo4j\" that have be...   \n",
       "9964  What is the average number of followers for us...   \n",
       "9965  Which users have a similarity score greater th...   \n",
       "9966  Find the users who follow \"Neo4j\" and have men...   \n",
       "\n",
       "                          type         database  \n",
       "0     Simple Retrieval Queries  recommendations  \n",
       "1     Simple Retrieval Queries  recommendations  \n",
       "2     Simple Retrieval Queries  recommendations  \n",
       "3     Simple Retrieval Queries  recommendations  \n",
       "4     Simple Retrieval Queries  recommendations  \n",
       "...                        ...              ...  \n",
       "9962        Multi-step Queries          twitter  \n",
       "9963        Multi-step Queries          twitter  \n",
       "9964        Multi-step Queries          twitter  \n",
       "9965        Multi-step Queries          twitter  \n",
       "9966        Multi-step Queries          twitter  \n",
       "\n",
       "[9967 rows x 3 columns]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.DataFrame.from_records(all_questions)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e3667e58-a621-4402-a9f4-d5f5b0323eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.drop_duplicates(subset='question').to_csv('gemini_questions.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f6ab397-a740-4fba-b5fa-332860feaf1a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
